import math
import time
import numpy as np
import pandas as pd
import coin.base.database.util as db_util
import google.protobuf.json_format as json_format

from absl import app, flags
from functools import reduce
from coin.proto.coin_query_pb2 import AccountBalance, AccountPosition
from coin.support.accounting.logic.db_util import (DBOutputEstimate, DBOutputPosition)
from coin.support.accounting.logic.format_util import (pretty_time_delta, move_df_nan_back,
                                                       format_currency_df, move_df_nan_back, print_total_impl)
from coin.support.accounting.logic.db_util import DBOutputPosition, DBOutputEstimate
from xunkemgmt.support.accounting.logic.query_request import _gen_engine
from coin.support.accounting.logic.memcached_querier import \
    MemcachedAccountInfoQuerier
from xunkemgmt_client.client.util.query_util import query_accounts

FLAGS = flags.FLAGS

flags.DEFINE_string('source_type', 'memcached', 'mysql, memcached')

flags.DEFINE_string('account_group', 'Xunke', 'account_group')

flags.DEFINE_string('memcached_config',
                    "../../coin_deploy/support_accounting/db_config/memcached_config_prod.json",
                    'path of database config.')


def _convert_estimate_to_df(estimates):
  estimates = [estimate for estimate in estimates if isinstance(estimate, DBOutputEstimate)]
  total_values, accys = [], []
  for estimate in estimates:
    estimate_value = estimate.balance_estimate or estimate.position_estimate
    total_values.append(estimate_value.estimate.total_value)
    accys.append(estimate_value.estimate.accounting_currency)
  time_to = [pretty_time_delta(time.time() - estimate.query_timestamp) for estimate in estimates]
  df = pd.DataFrame(estimates, columns=DBOutputEstimate._fields)
  df['time_to'] = pd.Series(time_to, index=df.index)
  df['total_value'] = pd.Series(total_values, index=df.index)
  df['accounting_currency'] = pd.Series(accys, index=df.index)
  return df


def _convert_position_to_df(positions):
  positions = [position for position in positions if isinstance(position, DBOutputPosition)]
  time_to = [pretty_time_delta(time.time() - position.query_timestamp) for position in positions]
  df = pd.DataFrame(positions, columns=DBOutputPosition._fields)
  df.rename(columns={'position': 'position_proto'}, inplace=True)
  df['time_to'] = pd.Series(time_to, index=df.index)
  return df


def get_balance_estimate(engine, account_data, source_type, memcached_config=None):
  if source_type == 'mysql':
    data2 = pd.read_sql(
        open("coin/support/accounting/query_sql/last_balance_estimate_per_account_currency.sql",
             'r').read(),
        engine)
    data2["time_to"] = data2["time_to"].map(lambda x: pretty_time_delta(x.total_seconds()))
  elif source_type == 'memcached':
    request_list = []
    business_units = list(account_data['business_unit'].unique())
    request_list = query_accounts(active=True, business_units=business_units, as_proto=True)
    db_config = db_util.read_db_config(memcached_config or FLAGS.memcached_config)
    querier = MemcachedAccountInfoQuerier(db_config=db_config)
    response_list = querier.query_balance_estimate_history(request_list)
    data2 = _convert_estimate_to_df(response_list)
  else:
    raise ValueError('Unknown source_type: %s', source_type)

  data = account_data.merge(data2, how="left", on="account_id")
  data["account_id_time_to"] = data["account_id"].map(str).fillna(
      "") + " | " + data["time_to"].fillna("__:__:__")
  data["accounting_currency"] = data["accounting_currency"].fillna("___")

  data['total_value'] = data.apply(lambda x: np.nan if x['account_group'] == 'HashedFund' and x[
      'accounting_currency'] != 'BTC' else x['total_value'],
                                   axis=1)

  def aggf(x):
    return x

  datat = pd.pivot_table(data,
                         index=['market_type', 'exchange', 'owner'],
                         columns='accounting_currency',
                         values='total_value',
                         aggfunc=aggf)
  datat = move_df_nan_back(datat, datat.columns)
  datat.loc[:, "VALID"] = reduce(
      lambda x, y: x | y, [~datat[c].isnull() for c in datat.columns])
  datat.loc['Total'] = datat.sum(axis=0, numeric_only=True)

  return datat


def get_position_estimate(engine, account_data, source_type):
  if source_type == 'mysql':
    data2 = pd.read_sql(
        open("coin/support/accounting/query_sql/last_position_estimate_per_account_currency.sql",
             'r').read(),
        engine)

    data2["time_to"] = data2["time_to"].map(lambda x: pretty_time_delta(x.total_seconds()))
  elif source_type == 'memcached':
    request_list = []
    business_units = list(account_data['business_unit'].unique())
    request_list = query_accounts(active=True, business_units=business_units)
    db_config = db_util.read_db_config(FLAGS.memcached_config)
    querier = MemcachedAccountInfoQuerier(db_config=db_config)
    response_list = querier.query_position_estimate_history(request_list)
    data2 = _convert_estimate_to_df(response_list)
  else:
    raise ValueError('Unknown source_type: %s', source_type)

  data = account_data.merge(data2, how="left", on="account_id")
  data["account_id_time_to"] = data["account_id"].map(str).fillna(
      "") + " | " + data["time_to"].fillna("__:__:__")
  data["accounting_currency"] = data["accounting_currency"].fillna("___")

  def aggf(x):
    return x

  datat = pd.pivot_table(data,
                         index=['market_type', 'exchange', 'owner'],
                         columns='accounting_currency',
                         values='total_value',
                         aggfunc=aggf)

  datat = move_df_nan_back(datat, datat.columns)
  datat.loc[:, "VALID"] = reduce(
      lambda x, y: x | y, [~datat[c].isnull() for c in datat.columns])

  datat.loc['Total'] = datat.sum(axis=0, numeric_only=True)
  datat.columns = [x[0] if isinstance(x, tuple) else x for x in datat.columns]

  return datat


def get_position_info(engine, account_data, source_type):
  def convert_proto(x):
    try:
      return json_format.Parse(x, AccountPosition())
    except Exception:
      return None

  if source_type == 'mysql':
    data2 = pd.read_sql(
        open("coin/support/accounting/query_sql/last_position_per_account.sql", 'r').read(), engine)
    data2["time_to"] = data2["time_to"].map(lambda x: pretty_time_delta(x.total_seconds()))
    data2["position_proto"] = data2["position_dict"].map(convert_proto)
  elif source_type == 'memcached':
    request_list = query_accounts(active=True, business_units=['Coin', 'CoinFOF'])
    db_config = db_util.read_db_config(FLAGS.memcached_config)
    querier = MemcachedAccountInfoQuerier(db_config=db_config)
    response_list = querier.query_account_position_history(request_list)
    data2 = _convert_position_to_df(response_list)
  else:
    raise ValueError('Unknown source_type: %s', source_type)

  def extract_liq_price(x):
    try:
      if len(x.each_position) > 0:
        if x.each_position[0].net_position > 0:
          return x.each_position[0].liquidation_price
        elif x.each_position[0].net_position < 0:
          return -1.0 * x.each_position[0].liquidation_price
        else:
          return None
      else:
        return None
    except Exception:
      return None

  data2["liq_price"] = data2["position_proto"].map(extract_liq_price)
  data2["liq_price"] = data2["liq_price"].map(lambda x: "{:,.2f}".format(x)
                                              if x is not None and not math.isnan(x) else "")

  data = account_data.merge(data2, how="left", on="account_id")
  data["account_id_time_to"] = data["account_id"].map(str).fillna(
      "") + " | " + data["time_to"].fillna("__:__:__")

  def aggf(x):
    return x

  datat = pd.pivot_table(data,
                         index=['market_type', 'exchange', 'owner'],
                         columns=[],
                         values='liq_price',
                         aggfunc=aggf)

  if 'liq_price' not in datat:
    datat.loc[:, 'liq_price'] = [np.nan] * len(datat.index)
  datat.loc[:, "VALID"] = ~datat['liq_price'].isnull()

  datat["liq_price"].fillna("", inplace=True)

  datat.loc['Total'] = np.nan

  return datat


def account(engine, account_group):
  print("\n\nACCOUNT")
  account_data = query_accounts(active=True, business_units=['Coin', 'CoinFOF'])
  account_data = account_data[account_data['account_group'].isin(account_group)]

  source_type = FLAGS.source_type
  datata = get_balance_estimate(engine, account_data, source_type)
  datatb = get_position_estimate(engine, account_data, source_type)
  datatc = get_position_info(engine, account_data, source_type)

  datata = format_currency_df(datata, print_zero=True)
  datatb = format_currency_df(datatb, print_zero=True)
  datatc = format_currency_df(datatc, print_zero=True)

  datatt = pd.concat([datata, datatb, datatc],
                     axis=1,
                     levels=["Balance", "Position", "PositionInfo"],
                     sort=False)
  datatt.replace(np.NAN, '', inplace=True)
  print(datatt)
  print_total_impl(datatt.loc['Total'].to_string())


def main(_):
  pd.set_option('display.width', 1000)
  pd.set_option('display.max_rows', None)
  pd.set_option('display.max_columns', None)

  engine = _gen_engine()
  account_groups = [elem.strip() for elem in FLAGS.account_group.split(',')]
  account(engine, account_groups)


if __name__ == '__main__':
  app.run(main)
